from openai import OpenAI
import os
import logging
from typing import List, Union

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TextEmbedding:
    """文本嵌入向量化工具"""
    
    def __init__(self):
        """初始化通义千问嵌入模型"""
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("请设置DASHSCOPE_API_KEY环境变量")

        self.client = OpenAI(
            api_key=self.api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        self.model = "text-embedding-v4"
        self.dimensions = 1024
        self.batch_size = 10

    def generate_embeddings(self, texts: Union[str, List[str]]) -> List[List[float]]:
        """生成文本嵌入向量"""
        # 处理单文本输入
        if isinstance(texts, str):
            texts = [texts]

        # 过滤空文本
        valid_texts = []
        valid_indices = []
        for idx, text in enumerate(texts):
            if text and text.strip():
                valid_texts.append(text.strip())
                valid_indices.append(idx)

        # 初始化结果列表
        embeddings = [[] for _ in range(len(texts))]

        # 为空文本填充零向量
        for idx in range(len(texts)):
            if idx not in valid_indices:
                embeddings[idx] = [0.0] * self.dimensions
                logger.warning(f"文本索引 {idx} 为空，返回零向量")

        if not valid_texts:
            return embeddings

        # 分批处理
        try:
            for i in range(0, len(valid_texts), self.batch_size):
                batch = valid_texts[i:i + self.batch_size]
                batch_indices = valid_indices[i:i + self.batch_size]

                response = self.client.embeddings.create(
                    model=self.model,
                    input=batch,
                    dimensions=self.dimensions,
                    encoding_format="float"
                )

                # 提取嵌入向量
                for data, idx in zip(response.data, batch_indices):
                    embeddings[idx] = data.embedding

                logger.info(f"已处理 {min(i + self.batch_size, len(valid_texts))}/{len(valid_texts)} 个文本")

            return embeddings

        except Exception as e:
            logger.error(f"生成嵌入向量失败: {str(e)}")
            raise

    def get_embedding_dimension(self) -> int:
        """获取嵌入向量维度"""
        return self.dimensions

    def is_valid_text(self, text: str) -> bool:
        """检查文本是否有效"""
        return text and text.strip()


def main():
    """主函数 - 演示文本嵌入功能"""
    try:
        embedder = TextEmbedding()

        # 测试单文本
        single_text = "如何在Boss直聘上优化企业招聘信息？"
        single_embedding = embedder.generate_embeddings(single_text)
        logger.info(f"单文本嵌入 - 维度: {len(single_embedding[0])}")

        # 测试多文本（包含空文本）
        texts = [
            "求职者如何提高简历的曝光率？",
            "企业如何筛选合适的候选人？",
            "",  # 空文本测试
            "面试邀约的最佳时间是什么时候？",
            "如何撰写吸引人的职位描述？"
        ]
        batch_embeddings = embedder.generate_embeddings(texts)
        logger.info(f"批量嵌入 - 总数量: {len(batch_embeddings)}")
        logger.info(f"各向量维度: {[len(emb) for emb in batch_embeddings]}")

    except Exception as e:
        logger.error(f"测试失败: {str(e)}")


if __name__ == "__main__":
    main()